import json
import time
import vtkmodules.all as vtk

from CloudEngine.core.func.parse_attributes import parse_attributes


def _float_arr_to_list(float_array: vtk.vtkFloatArray):
    tmp = []
    num_values = float_array.GetNumberOfValues()
    for i in range(num_values):
        tmp.append(float_array.GetValue(i))
    return tmp


class ICgnsReader:
    reader_instance = None
    series_reader_instance = None
    blocks = None
    attributes = None

    def __init__(self):
        self.reader_instance = vtk.vtkCGNSReader()
        self.series_reader_instance = vtk.vtkCGNSFileSeriesReader()
        self.blocks = []
        self.attributes = {}

    def simple_read(self, path: str):
        self.reader_instance.SetFileName(path)
        self.reader_instance.Update()

    def read(self, path: str):
        start = time.time()
        self.reader_instance.SetFileName(path)
        # ! 特喵的
        self.reader_instance.UpdateInformation()
        self.reader_instance.EnableAllBases()
        self.reader_instance.EnableAllCellArrays()
        self.reader_instance.EnableAllPointArrays()
        self.reader_instance.EnableAllFamilies()
        self.reader_instance.EnableAllFaceArrays()
        # 读取数据
        self.reader_instance.Update()
        multi_block = self.reader_instance.GetOutput()
        self.parse_block(multi_block)

        print("cgns解析消耗时间", time.time() - start, "秒")

    # 递归处理数据
    def parse_block(self, multi_block):
        for i in range(multi_block.GetNumberOfBlocks()):
            block = multi_block.GetBlock(i)
            if not block:
                continue

            # 🚨 调试用 误删
            # print(f"\n--- Block {i} ---")
            # print("Type:", block.GetClassName())

            # 判断是否是数据集类型（如 UnstructuredGrid、StructuredGrid 等）
            if isinstance(block, vtk.vtkDataSet):
                dataset = block
                print("数据集", dataset)

                # 🍪 表面抽取 如果要渲染体网格另外处理
                surface_filter = vtk.vtkDataSetSurfaceFilter()
                surface_filter.SetInputData(dataset)  # 输入为体积数据
                surface_filter.Update()
                poly_data = surface_filter.GetOutput()  # 输出为表面多边形数据

                # 临时的对象
                tmp_block = {
                    "name": f"Block {i}",
                    "points": [],
                    "cells": [],
                    "CellSize": [],
                    "attributes": [],
                }
                # 获取点数据（通常是物理量，如速度、压力）
                attrs = parse_attributes(poly_data)
                for key in attrs:
                    tmp_attr = {
                        "name": key,
                        "size": attrs[key].size,
                        "dx": _float_arr_to_list(attrs[key].dx),
                        "dy": [],
                        "dz": [],
                        "magtitude": [],
                    }
                    if attrs[key].size == 3:
                        tmp_attr["dy"] = _float_arr_to_list(attrs[key].dy)
                        tmp_attr["dz"] = _float_arr_to_list(attrs[key].dz)
                        tmp_attr["magtitude"] = _float_arr_to_list(attrs[key].magtitude)

                    tmp_block["attributes"].append(tmp_attr)

                # 获取点坐标
                points = poly_data.GetPoints()
                if points:
                    for i in range(points.GetNumberOfPoints()):
                        point = points.GetPoint(i)
                        # #️⃣ 注入所有顶点坐标的数据
                        tmp_block["points"].append(point[0])
                        tmp_block["points"].append(point[1])
                        tmp_block["points"].append(point[2])

                # 获取单元信息（拓扑结构）
                for i in range(poly_data.GetNumberOfCells()):
                    cell = poly_data.GetCell(i)
                    # print("cell 单元信息", cell)
                    tmp_block["CellSize"].append(cell.GetNumberOfPoints())
                    for i in range(cell.GetNumberOfPoints()):
                        point_id = cell.GetPointId(i)
                        # #️⃣ 注入所有单元的顶点索引
                        tmp_block["cells"].append(point_id)

                self.blocks.append(tmp_block)
            # 如果是 MultiBlock 的嵌套结构，递归处理（可选）
            elif isinstance(block, vtk.vtkMultiBlockDataSet):
                self.parse_block(block)

    def to_standard_json_structure(self):
        print("start 转换")
        start = time.time()
        tmp = {
            "blocks": [],
            "nodes": [],
            "nodeIds": [],
            "elements": [],
            "attributes": {},
        }
        point_flag = 0
        cell_flag = 0
        physical_flag = 0
        for block in self.blocks:
            # 用来标记获取单元内点ID的偏移量
            v_offset = 0
            tmp["nodes"] += block["points"]
            point_count: int = int(len(block["points"]) / 3)
            for i in range(point_count):
                tmp["nodeIds"].append(i + point_flag)
            cell_count: int = len(block["CellSize"])
            for i in range(cell_count):
                cell_size = block["CellSize"][i]
                node_ids = []
                for j in range(cell_size):
                    # 当前block的点索引 + 偏移量
                    node_ids.append(block["cells"][j + v_offset] + point_flag)
                v_offset += cell_size
                ele = {
                    "type": "CTRIA3" if cell_size == 3 else "CQUAD4",
                    "eid": i + cell_flag,
                    "nodeIds": node_ids,
                    "faceIds": [i + cell_flag],
                    "faceNums": [cell_size - 1],
                    "physicalTag": physical_flag,
                }
                tmp["elements"].append(ele)
            tmp["blocks"].append(
                {
                    "name": block["name"],
                    "elements": [cell_flag, cell_flag + cell_count],
                    "bid": physical_flag,
                }
            )

            block["bid"] = physical_flag
            point_flag += point_count
            cell_flag += cell_count
            physical_flag = physical_flag + 1

            # 处理属性数据
            if len(block["attributes"]) > 0:
                for obj in block["attributes"]:
                    if obj["size"] == 3:
                        if not self.attributes.get(obj["name"]):
                            self.attributes[obj["name"]] = {
                                "name": obj["name"],
                                "size": obj["size"],
                                "dx": obj["dx"],
                                "dy": obj["dy"],
                                "dz": obj["dz"],
                                "magtitude": obj["magtitude"],
                            }
                        else:
                            self.attributes[obj["name"]]["dx"].extend(obj["dx"])
                    else:
                        if not self.attributes.get(obj["name"]):
                            self.attributes[obj["name"]] = {
                                "name": obj["name"],
                                "size": obj["size"],
                                "dx": obj["dx"],
                            }
                        else:
                            self.attributes[obj["name"]]["dx"].extend(obj["dx"])
        tmp["attributes"] = self.attributes
        print("结束", time.time() - start, "单元数量", cell_flag)
        return tmp

    def to_json(self):
        return json.dumps(self.to_standard_json_structure())

    def get_attribute_list(self):
        print(self.attributes.keys())
